from fastapi import APIRouter, WebSocket, WebSocketDisconnect
import os
import json
import asyncio
import dashscope
import queue
import logging
from dashscope.audio.tts_v2 import *
from dashscope.audio.asr import *
from logger import logging

router = APIRouter()

# 初始化API密钥
try:
    from dotenv import load_dotenv
    load_dotenv()
    dashscope.api_key = os.getenv('API_KEY')
    logging.info("TTS模块API密钥初始化成功")
except Exception as e:
    logging.error(f"初始化API密钥失败: {e}")


class WSCallback():
    """WebSocket回调类，用于处理TTS流式输出"""

    def __init__(self, websocket: WebSocket, loop: asyncio.AbstractEventLoop):
        self.websocket = websocket
        self.loop = loop
        self.completed = queue.Queue()

    def on_open(self):
        pass

    def on_complete(self):
        self.completed.put(True)

    def on_error(self, message: str):
        logging.error(f'[tts]aliyun.websocket服务异常: {message}')
        self.completed.put(False)

    def on_close(self) -> None:
        pass

    def on_event(self, message) -> None:                    
        pass

    def on_data(self, data: bytes) -> None:
        def safe_send():
            try:
                task = asyncio.create_task(self.websocket.send_bytes(data))
                task.add_done_callback(self._handle_send_exception)
            except Exception as e:
                logging.error(f'[tts]创建发送任务失败：{e}')
        self.loop.call_soon_threadsafe(safe_send)

    def _handle_send_exception(self, task):
        try:
            task.result()
        except Exception as e:
            logging.error(f'[tts]WebSocket发送数据失败：{e}')


def check_sentence(text):
    """检查字符串是日文、中文还是韩文"""
    word_counts = {
        'japanese': 0,
        'chinese': 0,
        'korean': 0
    }
    for char in text:
        # 检查字符是否在日文相关范围内
        if ('\u3040' <= char <= '\u309F' or  # 平假名
                '\u30A0' <= char <= '\u30FF'):  # 片假名
            word_counts['japanese'] += 1
        elif ('\uAC00' <= char <= '\uD7AF'):  # 韩文
            word_counts['korean'] += 1
        else:
            word_counts['chinese'] += 1
    # 挑选单词数量最大的
    max_key = max(word_counts, key=word_counts.get)
    return max_key


async def tts_task(query: str, websocket: WebSocket):
    """处理LLM对话和TTS任务"""

    try:
        loop = asyncio.get_event_loop()
        synthesizer_callback = WSCallback(websocket=websocket, loop=loop)
        # 创建不同语言的语音合成器
        chinese_english_synthesizer = SpeechSynthesizer(
            model='cosyvoice-v2',
            voice=os.getenv('ZH_EN_VOICE'),
            format=AudioFormat.PCM_22050HZ_MONO_16BIT,
            callback=synthesizer_callback,
        )
        japanese_synthesizer = SpeechSynthesizer(
            model='cosyvoice-v2',
            voice=os.getenv('JA_VOICE'),
            format=AudioFormat.PCM_22050HZ_MONO_16BIT,
            callback=synthesizer_callback,
        )
        korean_synthesizer = SpeechSynthesizer(
            model='cosyvoice-v2',
            voice=os.getenv('KO_VOICE'),
            format=AudioFormat.PCM_22050HZ_MONO_16BIT,
            callback=synthesizer_callback,
        )
        sentence_type = check_sentence(query)
        # 根据语言类型选择合成器
        if sentence_type == 'japanese':
            japanese_synthesizer.streaming_call(query)
        elif sentence_type == 'korean':
            korean_synthesizer.streaming_call(query)
        else:
            chinese_english_synthesizer.streaming_call(query)
        await asyncio.sleep(0.05)
        if sentence_type == 'japanese':
            japanese_synthesizer.async_streaming_complete()
        elif sentence_type == 'korean':
            korean_synthesizer.async_streaming_complete()
        else:
            chinese_english_synthesizer.async_streaming_complete()
        # 等待TTS完成
        while True:
            if not synthesizer_callback.completed.empty():
                await websocket.send_text(json.dumps({'type': 'end'}))
                break
    except Exception as e:
        logging.error(f"TTS任务处理失败: {e}")
        await websocket.send_text(json.dumps({'type': 'error', 'message': str(e)}))


@router.websocket("/ws/tts")
async def websocket_tts(websocket: WebSocket):
    """WebSocket TTS端点"""
    await websocket.accept()
    logging.info("WebSocket连接已建立")

    try:
        while True:
            # 接收客户端消息
            data = await websocket.receive_text()
            logging.info(f"收到客户端消息: {data}")

            try:
                message = json.loads(data)
                if 'text' in message:
                    # 处理TTS消息
                    logging.info(f"接收到的消息: {message['text']}")
                    await tts_task(message['text'], websocket)
                else:
                    await websocket.send_text(json.dumps({
                        'type': 'error',
                        'message': '消息格式错误，需要包含text字段'
                    }))
            except json.JSONDecodeError:
                await websocket.send_text(json.dumps({
                    'type': 'error',
                    'message': '消息必须是有效的JSON格式'
                }))

    except WebSocketDisconnect:
        logging.info("WebSocket连接已断开")
    except Exception as e:
        logging.error(f"WebSocket错误: {e}")
        try:
            await websocket.send_text(json.dumps({
                'type': 'error',
                'message': f'服务器错误: {str(e)}'
            }))
        except:
            pass
